import com.sun.jna.Pointer;
import library.WintunLibrary;
import org.pcap4j.packet.*;
import org.pcap4j.packet.namednumber.IpNumber;
import org.pcap4j.packet.namednumber.IpVersion;
import org.pcap4j.packet.namednumber.TcpPort;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.*;

import java.util.Arrays;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class TcpData {

    static ConcurrentHashMap<String, TcpPipe> pipeMap = new ConcurrentHashMap<>();

    public static void parse(IpPacket packet, Pointer pointer) {

        //通过源ip和源端口来判断同一个tcp流
        String key = getKey(packet);

        TcpPipe tcpPipe = pipeMap.get(key);
        if (tcpPipe == null) {
            //创建一个新的对象
            tcpPipe = new TcpPipe(pointer);
            //防止多线程覆盖对象
            pipeMap.putIfAbsent(key, tcpPipe);

        }
        tcpPipe.receive(packet);

    }

    private static String getKey(IpPacket packet) {
        InetAddress srcAddr = packet.getHeader().getSrcAddr();
        TcpPacket tcpPacket = (TcpPacket) packet.getPayload();
        TcpPort srcPort = tcpPacket.getHeader().getSrcPort();
        return srcAddr + ":" + srcPort.valueAsInt();
    }

    private static void remove(IpPacket packet) {
        pipeMap.remove(getKey(packet));
    }


    static class TcpPipe {

        public Pointer pointer;

        private AtomicBoolean synState = new AtomicBoolean(false);
        private AtomicLong replySeq = new AtomicLong(1000);
        private AtomicLong replyAck = new AtomicLong(0);

        private AtomicBoolean sendFin = new AtomicBoolean(false);
        private AtomicBoolean replyFin = new AtomicBoolean(false);


        private IpPacket srcIpPacket;

        private OutputStream outputStream;

        private InputStream inputStream;

        private Socket socket;

        private Thread readThread;

        public TcpPipe(Pointer pointer) {
            this.pointer = pointer;
        }

        private void writeData(byte[] data) {

            Pointer packetPointer = WintunLibrary.INSTANCE.WintunAllocateSendPacket(this.pointer, data.length);

            if (packetPointer != null) {
                // 构造 ICMP 数据包
                packetPointer.write(0, data, 0, data.length);

                // 发送数据包
                WintunLibrary.INSTANCE.WintunSendPacket(this.pointer, packetPointer);

            } else {
                System.err.println("Failed to allocate send packet.");
            }
        }

        public void receive(IpPacket packet) {

            TcpPacket tcpPacket = (TcpPacket) packet.getPayload();

            boolean syn = tcpPacket.getHeader().getSyn();
            boolean fin = tcpPacket.getHeader().getFin();
            boolean psh = tcpPacket.getHeader().getPsh();
            boolean ack = tcpPacket.getHeader().getAck();
            boolean rst = tcpPacket.getHeader().getRst();

            InetAddress dstAddr = packet.getHeader().getDstAddr();
            int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();


            //只拦截 example.com
            if (dstAddr.getHostAddress().equals("172.29.1.25")) {
                try {
                    dstAddr = InetAddress.getByName("2606:2800:220:1:248:1893:25c8:1946");
                } catch (UnknownHostException e) {
                    throw new RuntimeException(e);
                }
            } else {
                close();
                return;
            }


            //认为结束
            if (sendFin.get() && replyFin.get()) {
                close();
                return;
            }

            if (syn) {
                //防止超时重发syn数据包
                if (!synState.getAndSet(true)) {

                    //一次握手
                    try {
                        socket = new Socket(dstAddr, dstPort);
                        //初始化一些参数
                        srcIpPacket = packet;
                        long seq = tcpPacket.getHeader().getSequenceNumberAsLong();
                        replyAck.set(seq + 1);


                        //说明连接成功,回复ack(第二次握手)
                        replySyn();

                        outputStream = socket.getOutputStream();

                        inputStream = socket.getInputStream();

                        //开启一个线程去读取服务端返回数据

                        readThread = new Thread(() -> {
                            try {
                                byte[] bt = new byte[1024];
                                int len = -1;
                                while (!Thread.interrupted()) {
                                    if ((len = inputStream.read(bt)) != -1) {
                                        if (len != bt.length) {
                                            bt = Arrays.copyOf(bt, len);
                                        }
                                        //服务端回复数据包给客户端
                                        replyPsh(bt);
                                    }

                                    if (len == -1) {
                                        //发送Fin
                                        if (!sendFin.getAndSet(true)) {
                                            replyFin();
                                        }
                                        close();
                                        break;
                                    }

                                }
                            } catch (IOException e) {
                                e.printStackTrace();
                            }


                        });

                        readThread.setName(dstAddr + ":" + dstPort);
                        readThread.start();

                    } catch (IOException e) {
                        e.printStackTrace();
                        close();
                    }

                }

            } else {

                try {
                    Packet payload = tcpPacket.getPayload();

                    //有数据就发送到服务端
                    if (payload != null && payload.getRawData().length > 0) {

                        if (socket != null && socket.isConnected()) {
                            outputStream.write(payload.getRawData());
                            replyAck.addAndGet(payload.getRawData().length);
                            replyAck();
                        }

                    }

                    if (fin) {

                        replyFin.set(true);
                        replyAck.incrementAndGet();

                        replyAck(tcpPacket.getHeader().getAcknowledgmentNumber());

                        if (!sendFin.getAndSet(true)) {
                            replyFin();
                        }


                    } else if (psh) {

                    }

                    if (ack) {

                    }

                } catch (Exception e) {
                    e.printStackTrace();
                    close();
                }

            }

        }

        private void close() {
            if (readThread != null) {
                readThread.interrupt();
            }
            if (outputStream != null) {
                try {
                    outputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (inputStream != null) {
                try {
                    outputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (socket != null) {
                try {
                    socket.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (srcIpPacket != null) {
                remove(srcIpPacket);
            }

        }

        public void replySyn() {

            IpPacket newIpacket = buildIpPack(true, false, false, true, null);
            byte[] rawData = newIpacket.getRawData();
            writeData(rawData);

            replySeq.incrementAndGet();
        }

        public void replyFin() {

            long seq = replySeq.get();

            IpPacket newIpacket = buildIpPack(false, false, true, true, null, seq);
            byte[] rawData = newIpacket.getRawData();
            writeData(rawData);


            replySeq.set(seq + 1);
        }

        public void replyAck() {


            IpPacket newIpacket = buildIpPack(false, false, false, true, null);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);
        }

        public void replyPsh(byte[] data) {

            IpPacket newIpacket = buildIpPack(false, true, false, true, data);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);

            this.replySeq.getAndAdd(data.length);

        }


        public void replyAck(long seq) {


            IpPacket newIpacket = buildIpPack(false, false, false, true, null, seq);
            byte[] rawData = newIpacket.getRawData();

            writeData(rawData);
        }


        private IpPacket buildIpPack(boolean syn, boolean psh, boolean fin, boolean ack, byte[] data) {
            long seq = replySeq.get();
            return buildIpPack(syn, psh, fin, ack, data, seq);
        }

        private IpPacket buildIpPack(boolean syn, boolean psh, boolean fin, boolean ack, byte[] data, long seq) {

            TcpPacket localtcpPacket = (TcpPacket) srcIpPacket.getPayload();
            TcpPort srcPort = localtcpPacket.getHeader().getSrcPort();
            TcpPort dstPort = localtcpPacket.getHeader().getDstPort();


            InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
            InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();

            TcpPacket.Builder tcpBuilder;
            tcpBuilder = new TcpPacket.Builder()
                    .srcPort(dstPort)
                    .dstPort(srcPort)
                    .sequenceNumber((int) seq)
                    .window(Short.MIN_VALUE)
                    .acknowledgmentNumber((int) this.replyAck.get())
                    .syn(syn)
                    .psh(psh)
                    .fin(fin)
                    .ack(ack)
                    .srcAddr(dstAddr)
                    .dstAddr(srcAddr)
                    .correctChecksumAtBuild(true)
                    .correctLengthAtBuild(true);


            if (psh) {
                UnknownPacket.Builder unknownBuilder = new UnknownPacket.Builder();
                unknownBuilder.rawData(data);
                tcpBuilder.payloadBuilder(unknownBuilder);
            }


            //创建一个ip包
            IpPacket newIpPacket;
            if (srcIpPacket instanceof IpV6Packet) {
                newIpPacket = new IpV6Packet.Builder()
                        .version(IpVersion.IPV6)
                        .srcAddr((Inet6Address) dstAddr)
                        .dstAddr((Inet6Address) srcAddr)
                        .correctLengthAtBuild(true)
                        .payloadBuilder(tcpBuilder)
                        .trafficClass(IpV6SimpleTrafficClass.newInstance((byte) 0))
                        .flowLabel(IpV6SimpleFlowLabel.newInstance((byte) 0))
                        .nextHeader(IpNumber.IPV6)
                        .build();
            } else {

                newIpPacket = new IpV4Packet.Builder()
                        .version(IpVersion.IPV4)
                        .ttl((byte) 45)
                        .tos(IpV4Rfc1349Tos.newInstance((byte) 0))
                        .protocol(IpNumber.TCP)
                        .srcAddr((Inet4Address) dstAddr)
                        .dstAddr((Inet4Address) srcAddr)
                        .correctLengthAtBuild(true)
                        .correctChecksumAtBuild(true)
                        .payloadBuilder(tcpBuilder)
                        .identification((byte) getIpId())
                        .build();
            }

            return newIpPacket;
        }
    }

    private static AtomicInteger ipId = new AtomicInteger(0);

    public static int getIpId() {
        return ipId.incrementAndGet();
    }


}
